# -*- coding: utf-8 -*-
# May Saint Diana bless your coding!
# Created Time : 2021/12/2 21:50
# Author:Zhou Chenghao

import numpy as np
import torch
from torch.autograd import Variable
import random
import os
import matplotlib.pyplot as plt
from ANC_tools import ANR, insertToHead, insertToTail
import torch.nn as nn

# 读取实验数据
# path = "次级通道建模数据.csv"  # 输入是传递给次级扬声器播放的信号（高斯白噪声，能激发出所有频段），输出是在误差传感器处采集到的误差信号
path = "非线性3.csv"  # 一个参考信号和该信号经过非线性过程的对应输出
data = np.loadtxt(path, delimiter=",")

# 构造简单训练数据
x = []
y = []

for i in range(1000):
    time_point = random.randint(31, 45000)  # 找一个随机时刻
    x.append(data[time_point - 31:time_point + 1, 0])  # 训练输入数据是当前及过去15个时刻的输出值
    y.append(data[time_point, 1])  # 输出是当前时刻的传感器数据  认为：当前的y和过去若干时刻的输入有关
# x y 放进numpy，再转成pytorch的tensor数据
x = np.array(x)
y = np.array(y)
x = Variable(torch.from_numpy(x)).to(torch.float)
y = Variable(torch.from_numpy(y[:, np.newaxis])).to(torch.float)
# 训练数据：1000组，每组输入是32*1，输出是8*1，即：构建当前时刻的中间信号和过去若干时刻（32）的次级扬声器输出信号之间的映射关系:
# y(n),y(n-1),...,y(n-7)=f[x(n),x(n-1),x(n-2),x(n-3),...,x(n-15)]

#  自定义一个全连接网络 32*64*32*8  当前网络的输出和之前32个时间点的输入有关
net = torch.nn.Sequential(
    torch.nn.Linear(32, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 32),
    torch.nn.ReLU(),
    torch.nn.Linear(32, 8))

# 如果存在训练过的网络参数文件，直接读取
if os.path.exists('bp_anc_net_params.pkl'):
    net.load_state_dict(torch.load('bp_anc_net_params.pkl'))  # 提取net的参数

# secondaryPathModel = torch.tensor([[0, 0, 0, 1, 0.5, 0.2, -0.1, 0.08]])
secondaryPathModel = torch.tensor([[0.08, -0.1, 0.2, 0.5, 1, 0, 0, 0]])  # 最前面的是最早之前的 y(n-7)


class My_loss(nn.Module):
    """
    自定义损失函数
    """

    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        """
        网络得到的输出要继续经过次级路径，得到降噪点处的信号
        :param x:
        :param y:
        :return:
        """
        filtered_y = torch.matmul(x, secondaryPathModel.t())
        return torch.pow(torch.mean(torch.abs(filtered_y + y)), 2)


optimizer = torch.optim.SGD(net.parameters(), lr=0.1)  # 用SGD优化器优化，lr为学习率，将net的参数传入
loss_func = My_loss()

plt.ion()  # 画图
plt.show()


print("----------开始BP训练----------")
for t in range(1000):
    # 通过网络计算得到的预测值  暂时没有用批训练，因此每次都是将1000个x先后输入网络，得到1000个输出y(n),y(n-1),y(n-2)...y(n-7)
    prediction = net(x)
    # 计算y'和y的误差，得到损失函数
    loss = loss_func(prediction, y)
    optimizer.zero_grad()  # 将net.parameters中上一批梯度清零
    loss.backward()  # 反向传递  计算节点梯度
    optimizer.step()  # 优化节点
    # print(loss.data, loss.grad, loss.grad_fn) # 展示梯度

    if t % 5 == 0:  # 动态显示
        # plot and show learning process
        plt.cla()
        plt.plot(y.data.numpy(), lw=5)
        filtered_output = np.dot(prediction.data.numpy(), secondaryPathModel.t().data.numpy())
        plt.plot(filtered_output, 'r-', lw=5)
        plt.text(25, 3, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color': 'green'})
        plt.pause(0.1)

print("----------训练完成！----------")
# 保存网络的参数
# torch.save(net.state_dict(), 'bp_anc_net_params.pkl')  # 只保存网络中的参数 (速度快, 占内存少)

plt.ioff()
plt.show()

# 验证拟合效果
signal_x = data[:, 0]
signal_y = data[:, 1]
cx = np.zeros(32)
BP_error = np.zeros(len(signal_x))
print("----------验证拟合效果----------")
for i in range(len(signal_x)):
    cx = insertToTail(cx, signal_x[i])  # 需要和网络的数据对应 从上至下应该是 x(n-15) 到 x(n)，所以要插入到末尾
    BP_x = Variable(torch.from_numpy(cx)).to(torch.float)
    BP_pred = net(BP_x)
    filtered_output = np.dot(BP_pred.data.numpy(), secondaryPathModel.t().data.numpy())
    BP_error[i] = filtered_output + signal_y[i]
print("平均误差：{0},最大误差：{1}".format(np.average(np.abs(BP_error)),
                                 np.max(np.abs(BP_error))))

plt.subplot(121)
plt.plot(signal_y)
plt.plot(BP_error)
plt.legend(['y', 'error'])
plt.title('BP')
plt.subplot(122)
plt.plot(ANR(signal_y, BP_error))
plt.title('BP_ANR')
plt.show()

"""
问题和思考：
1、虽然在仿真里面实现了不错的效果，但还是很难用于实际，原因如下：
    网络的输出是个8*1的，我们假想：网络能够给出y(n-7)...y(n)和x(n-31)...x(n)之间的关系
    但实际上，观察验证时各个时刻的输出，并没有发现延迟关系，所以网络实现的仅仅是数学上的拟合，物理上并没有达到目的
    因此，就无法确定某个时刻真正要发送给次级扬声器的控制信号（一个值，不是一串值），只能仿真，物理上无法实现
    《基于BP神经网络的车内有源噪声控制技术研究_李燚勇》这篇论文里面，是用BP模拟了传感器采集到的物理信号（路面、风噪）和降噪点处信号之间的
    非线性关系，再将网络输出交给滤波器实现ANC，并不是利用BP直接控制
"""
# TODO 做一下批处理
